import os
import numpy
import arcpy

import stormwaterutilities
import accumulatetonetwork
import pollutantdata as pd
from stormwaterutilities import addMessage

def CalculateScoresAndConcentrations(table, rainfallDuration):
    #Calculate net concentrations and magnitude scores
    inArray = arcpy.da.FeatureClassToNumPyArray(table, (['OID@', 'flow_cfs', 'TSS_lbs', 'PbD_lbs', 'TP_lbs', 'BOD_lbs', 'ecoli_cfu']))

    flow = inArray['flow_cfs']
    volume = flow * rainfallDuration * 24 * 60 * 60

    tssLoad = inArray['TSS_lbs']
    pbLoad = inArray['PbD_lbs']
    tpLoad = inArray['TP_lbs']
    bodLoad = inArray['BOD_lbs']
    ecoliLoad = inArray['ecoli_cfu']

    tssConcentrationNet = numpy.where(volume > 0, tssLoad / volume, 0) * 16018.4
    tssScoreNet = stormwaterutilities.mapValues(tssConcentrationNet, pd.tssScoreBasis, pd.scores)

    pbConcentrationNet = numpy.where(volume > 0, pbLoad / volume, 0) * 16018400
    pbScoreNet = stormwaterutilities.mapValues(pbConcentrationNet, pd.pbScoreBasis,  pd.scores)

    tpConcentrationNet = numpy.where(volume > 0, tpLoad / volume, 0) * 16018.4
    tpScoreNet = stormwaterutilities.mapValues(tpConcentrationNet, pd.tpScoreBasis,  pd.scores)

    bodConcentrationNet = numpy.where(volume > 0, bodLoad / volume, 0) * 16018.4
    bodScoreNet = stormwaterutilities.mapValues(bodConcentrationNet, pd.bodScoreBasis,  pd.scores)

    ecoliConcentrationNet = numpy.where(volume > 0, ecoliLoad / volume, 0) * 0.00353146667
    ecoliScoreNet = stormwaterutilities.mapValues(ecoliConcentrationNet, pd.ecoliScoreBasis,  pd.scores)

    numpy.seterr(divide='warn')

    #For now, magnitude score is unweighted mean of pollutants
    magnitudeScoreNet = (tssScoreNet + pbScoreNet + tpScoreNet + bodScoreNet + ecoliScoreNet) / 5

    addMessage('Writing net concentrations and scores to NumPy array')
    typeList = numpy.dtype([('oid', numpy.int32), ('volume_cu_ft', numpy.float32),
                            ('TSS_mg_per_L', numpy.float32), ('TSS_score', numpy.float32),
                            ('PbD_ug_per_L', numpy.float32), ('PbD_score', numpy.float32),
                            ('TP_mg_per_L', numpy.float32), ('TP_score', numpy.float32),
                            ('BOD_mg_per_L', numpy.float32), ('BOD_score', numpy.float32),
                            ('ecoli_cfu_per_100mL', numpy.float32), ('ecoli_score', numpy.float32),
                            ('magnitude_score', numpy.float32)])

    outArray = numpy.asarray(zip(inArray['OID@'], volume,
                                 tssConcentrationNet, tssScoreNet, pbConcentrationNet, pbScoreNet, tpConcentrationNet, tpScoreNet,
                                 bodConcentrationNet, bodScoreNet, ecoliConcentrationNet, ecoliScoreNet, magnitudeScoreNet), typeList)

    addMessage('Saving calculations back to feature class')
    arcpy.da.ExtendTable(table, 'OID@', outArray, 'oid', False)

    arcpy.AddField_management(table, 'wq_risk_score', 'FLOAT')
    arcpy.CalculateField_management(table, 'wq_risk_score', "!wq_susc_score! * !magnitude_score!", 'PYTHON')

def CreateBaseConcentrationRaster(landuseRasterPath, outRasterPath, pollutantConcentrationDictionary):

    landuseRaster = arcpy.sa.Raster(landuseRasterPath)
    addMessage('Picking mean pollutant concentrations from landuse raster.')

    temp = arcpy.env.workspace
    arcpy.env.workspace = r'd:\temp\swsp_working.gdb'

    #This method relies on the landuse codes being mapped to the 'Value' field of the landuse raster such that
    #RES:1, COM:2, IND:3, POS:4, VAC:5, TRU:6, TRL:7, TRM:8, TRH:9, WAT:10
    #We could make this more generic but it's simple and it works for now.
    concentrationRaster = arcpy.sa.Pick(landuseRaster, [
        pollutantConcentrationDictionary['RES'].mean,
        pollutantConcentrationDictionary['COM'].mean,
        pollutantConcentrationDictionary['IND'].mean,
        pollutantConcentrationDictionary['POS'].mean,
        pollutantConcentrationDictionary['VAC'].mean,
        pollutantConcentrationDictionary['TRU'].mean,
        pollutantConcentrationDictionary['TRL'].mean,
        pollutantConcentrationDictionary['TRM'].mean,
        pollutantConcentrationDictionary['TRH'].mean,
        0]) #Landuse 'Value' of 7 corresponds to 'WAT', which always has a load of 0

    addMessage('Saving concentration raster: ' + outRasterPath)
    concentrationRaster.save(outRasterPath)

    arcpy.env.workspace = temp

    return

def CreateNetConcentrationRasters(outputFolder): #TODO: Is this the right name for this method?
    runoff = r"\\ws18325\Stormwater D\RiskModel\Surface\LandClassification.gdb\RunoffCoefficient"
    input_table = r"\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality\PollutantData.gdb\bmp_effluent_concentrations"
    input_BMP = r"\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality\PollutantData.gdb\private_BMP"

    tempFolder = r"D:\Temp\swsp_working"
    if not os.path.exists(tempFolder): os.makedirs(tempFolder)

    database = os.path.join(tempFolder, "working.gdb")
    if arcpy.Exists(database):
        addMessage("Working dB exists")
    else:
        addMessage("Creating working dB")
        arcpy.CreateFileGDB_management(tempFolder, "working.gdb","CURRENT")

    #sde connection saved locally so that it is standardized
    # Connect to CGIS hub
    script_dir = os.path.dirname(os.path.abspath(__file__))
    sde_path = os.path.join(script_dir, "GISDB1.ROSE.PORTLAND.LOCAL.sde")

    taxlots =r"EGH_PUBLIC.ARCMAP_ADMIN.taxlots_pdx"
    egh_taxlots = os.path.join(sde_path, taxlots)

    #Pivot table to group by ACWA general type
    addMessage("Creating Pivot Table")
    pivot = os.path.join(database,"pivot_table")
    pivot_table = arcpy.PivotTable_management(input_table, 'gen_id', 'short_name', 'effluent_concentration_mean', pivot)

    #join fields from pollutant table to the private BMP points
    addMessage("Joining pollutant info to private facilities")
    BMP_layer = arcpy.MakeFeatureLayer_management(input_BMP, "BMP_layer")
    arcpy.AddJoin_management(BMP_layer, 'BMP_TYPE_GEN_ID', pivot_table, 'gen_id')

    JoinOutFile = os.path.join(database, "BMP_values_join")
    arcpy.CopyFeatures_management(BMP_layer, JoinOutFile)

    #spatial join facility values to associated taxlots
    addMessage("Spatial joining values to taxlots")
    out_sj = r'in_memory\lot_private_bmp_sj' # os.path.join(database,"lot_privateBMP_sj")
    arcpy.SpatialJoin_analysis(egh_taxlots, JoinOutFile, out_sj, "JOIN_ONE_TO_ONE", "KEEP_COMMON",
        "pivot_table_BOD \"BOD\" true true false 8 Double 0 0 ,Min,#,d:\\Temp\\swsp_working\\working.gdb\\BMP_values_join,pivot_table_BOD,-1,-1;pivot_table_ecoli \"ecoli\" true true false 8 Double 0 0 ,Min,#,d:\\Temp\\swsp_working\\working.gdb\\BMP_values_join,pivot_table_ecoli,-1,-1;pivot_table_PbD \"PbD\" true true false 8 Double 0 0 ,Min,#,d:\\Temp\\swsp_working\\working.gdb\\BMP_values_join,pivot_table_PbD,-1,-1;pivot_table_TP \"TP\" true true false 8 Double 0 0 ,Min,#,d:\\Temp\\swsp_working\\working.gdb\\BMP_values_join,pivot_table_TP,-1,-1;pivot_table_TSS \"TSS\" true true false 8 Double 0 0 ,Min,#,d:\\Temp\\swsp_working\\working.gdb\\BMP_values_join,pivot_table_TSS,-1,-1;pivot_table_Volume_Reduction \"Volume Reduction\" true true false 8 Double 0 0 ,Max,#,d:\\Temp\\swsp_working\\working.gdb\\BMP_values_join,pivot_table_Volume_Reduction,-1,-1",
        "INTERSECT", "", "")

    #Handle volume reduction
    addMessage("Saving volume reduction effectiveness to raster")

    bmpFlowControlRasterPath = os.path.join(outputFolder, "bmp_flow_control_effectiveness")
    #flowControlPath = r"in_memory\flowControlEff"
    arcpy.FeatureToRaster_conversion(out_sj, "pivot_table_Volume_Reduction", bmpFlowControlRasterPath, "3")
    bmpFlowControlRaster = arcpy.sa.Raster(bmpFlowControlRasterPath)

    flowControlRaster = arcpy.sa.Con(arcpy.sa.IsNull(bmpFlowControlRaster), 0, bmpFlowControlRaster / 100.0)
    flowControlRasterPath = os.path.join(outputFolder, "flow_control_effectiveness")
    flowControlRaster.save(flowControlRasterPath)
    #arcpy.Delete_management(flowControlPath)

    #calc the net runoff coefficient"
    addMessage("Creating the Net Runoff Coefficient raster")
    netRunoffCoeffRaster = arcpy.Raster(runoff) - (arcpy.Raster(runoff) * arcpy.Raster((flowControlRasterPath)))
    out = os.path.join(outputFolder, "net_runoff_coeff")
    netRunoffCoeffRaster.save(out)

    #Handle pollutants
    for pollutant in pd.pollutantList:
        addMessage("Saving net pollutant concentration raster: " + pollutant)
        pivotFieldName = 'pivot_table_' + pollutant
        outTempPath = r"in_memory\temp"
        arcpy.FeatureToRaster_conversion(out_sj, pivotFieldName, outTempPath, 3)

        netConcRasterPath = os.path.join(outputFolder, 'net_conc_' + pollutant)
        baseConcRasterPath = os.path.join(outputFolder, 'base_conc_' + pollutant)

        outRaster = arcpy.sa.Con(arcpy.sa.IsNull(outTempPath), baseConcRasterPath, outTempPath)

        outRaster.save(netConcRasterPath)
        arcpy.Delete_management(outTempPath)

def CreateNetLoadRaster(outputFolder, pollutant, rainfallDepth, gridCellArea, conversionFactor, netRunoffCoeff):
    netConcRasterPath = os.path.join(outputFolder, 'net_conc_' + pollutant)
    baseConcRasterPath = os.path.join(outputFolder, 'base_conc_' + pollutant)
    netLoadRasterPath = os.path.join(outputFolder, 'net_load_' + pollutant)

    addMessage("Creating net load raster from concentration raster: " + netLoadRasterPath)

    netConcRaster = arcpy.Raster(netConcRasterPath)
    baseConcRaster = arcpy.Raster(baseConcRasterPath)

    loadRaster = arcpy.sa.Con(netConcRaster < baseConcRaster, netConcRaster, baseConcRaster) * rainfallDepth * gridCellArea * conversionFactor * netRunoffCoeff
    loadRaster.save(netLoadRasterPath)

def CreateLoadRasters(links, nodes, catchments, flowDir, landuseRasterPath, outputFolder):
    CreateBaseConcentrationRaster(landuseRasterPath, outputFolder + r'\base_conc_tss',
        pd.tssLuConcentration)

    CreateBaseConcentrationRaster(landuseRasterPath, outputFolder + r'\base_conc_pbd',
        pd.pbLuConcentration)

    CreateBaseConcentrationRaster(landuseRasterPath, outputFolder + r'\base_conc_tp',
        pd.tpLuConcentration)

    CreateBaseConcentrationRaster(landuseRasterPath, outputFolder + r'\base_conc_bod',
        pd.bodLuConcentration)

    CreateBaseConcentrationRaster(landuseRasterPath, outputFolder + r'\base_conc_ecoli',
        pd.ecoliLuConcentration)

    CreateNetConcentrationRasters(outputFolder)

    netRunoffCoeff = os.path.join(outputFolder, 'net_runoff_coeff')

    CreateNetLoadRaster(outputFolder, 'tss', pd.annualRainfallDepth, 9, pd.mg_per_L_conversion, netRunoffCoeff)
    CreateNetLoadRaster(outputFolder, 'pbd', pd.annualRainfallDepth, 9, pd.ug_per_L_conversion, netRunoffCoeff)
    CreateNetLoadRaster(outputFolder, 'tp', pd.annualRainfallDepth, 9, pd.mg_per_L_conversion, netRunoffCoeff)
    CreateNetLoadRaster(outputFolder, 'bod', pd.annualRainfallDepth, 9, pd.mg_per_L_conversion, netRunoffCoeff)
    CreateNetLoadRaster(outputFolder, 'ecoli', pd.annualRainfallDepth, 9, pd.cfu_per_100_mL_conversion, netRunoffCoeff)

    netRunoffCoeff = os.path.join(outputFolder, 'net_runoff_coeff')
    (netConcRasters, baseConcRasters, netLoadRasters) = ([],  [], [])

    for pollutant in pd.pollutantList:
        netConcRasters.append(os.path.join(outputFolder, "net_conc_" + pollutant))
        baseConcRasters.append(os.path.join(outputFolder, "base_conc_" + pollutant))
        netLoadRasters.append(os.path.join(outputFolder, "net_load_" + pollutant))

    #Accumulate net runoff generating area
    addMessage('Accumulating net volume.')
    accumulatetonetwork.AccumulateToNetwork(links, nodes, catchments, flowDir,
            netRunoffCoeff, 'accumulated_area_sq_ft', 9)

    #Accumulate pollutant load
    for pollutant, netLoadRaster in zip(pd.pollutantList, netLoadRasters):
        addMessage('Accumulating %s'%netLoadRaster)
        accumulatetonetwork.AccumulateToNetwork(links, nodes, catchments, flowDir,
            netLoadRaster, pollutant + '_' + pd.pollutantLoadUnits[pollutant], 1)

def Decumulate(links, nodes, catchments):
    wqDir = r'\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality'
    pollutantGdb = os.path.join(wqDir, 'PollutantData.gdb')
    publicBmps = os.path.join(pollutantGdb, 'public_BMP')

    stormwaterutilities.CopyFieldFromFeatureClass(nodes, 'UNITID', 'prf_type_id', publicBmps, 'Nearest_Hansen', 'PRF_GEN_TYPE_ID')

    decum_args = []
    decum_args.append('Decumulate')
    decum_args.append(links)
    decum_args.append(nodes)
    decum_args.append(catchments)
    decum_args.append(os.path.join(pollutantGdb, 'bmp_effluent_concentrations'))

    # TSS
    decum_args.append('TSS')
    decum_args.append('flow_cfs')
    decum_args.append('28.3168') # cu_ft to L
    decum_args.append('TSS_lbs')
    decum_args.append('453592.37') # lbs to mg
    decum_args.append('flow_cfs_temp')
    decum_args.append('TSS_lbs')

    stormwaterutilities.callGpTool(decum_args)

    # Pb, d
    decum_args[5] = 'PbD'
    decum_args[6] = 'flow_cfs'
    decum_args[7] = '28.3168' # cu_ft to L
    decum_args[8] = 'PbD_lbs'
    decum_args[9] = '453592370' # lbs to ug (micrograms)
    decum_args[10] = 'flow_cfs_temp'
    decum_args[11] = 'PbD_lbs'

    stormwaterutilities.callGpTool(decum_args)

    # TP
    decum_args[5] = 'TP'
    decum_args[6] = 'flow_cfs'
    decum_args[7] = '28.3168' # cu_ft to L
    decum_args[8] = 'TP_lbs'
    decum_args[9] = '453592.37' # lbs to mg
    decum_args[10] = 'flow_cfs_temp'
    decum_args[11] = 'TP_lbs'

    stormwaterutilities.callGpTool(decum_args)

    # BOD
    decum_args[5] = 'BOD'
    decum_args[6] = 'flow_cfs'
    decum_args[7] = '28.3168' # cu_ft to L
    decum_args[8] = 'BOD_lbs'
    decum_args[9] = '453592.37' # lbs to mg
    decum_args[10] = 'flow_cfs_temp'
    decum_args[11] = 'BOD_lbs'

    stormwaterutilities.callGpTool(decum_args)

    # E. coli
    decum_args[5] = 'ecoli'
    decum_args[6] = 'flow_cfs'
    decum_args[7] = '283.168' # cu_ft to 100 mL
    decum_args[8] = 'ecoli_cfu'
    decum_args[9] = '1.0' # CFU to MPN (assuming no difference at this point)
    decum_args[10] = 'flow_cfs' # Only the last decumulate call should overwrite the actual flow_cfs field
    decum_args[11] = 'ecoli_cfu'

    stormwaterutilities.callGpTool(decum_args)

    arcpy.DeleteField_management(links, 'flow_cfs_temp')
    arcpy.DeleteField_management(nodes, 'flow_cfs_temp')


def CalculateBaseFlowAndVolume(table, rainfallDepth, rainfallDuration):
    inArray = arcpy.da.FeatureClassToNumPyArray(table, (['OID@', 'accumulated_area_sq_ft']))

    volume = inArray['accumulated_area_sq_ft'] * rainfallDepth / 12
    flow = volume / (rainfallDuration * 24 * 60 * 60)

    addMessage('Writing volume and flow to NumPy array')
    typeList = numpy.dtype([('oid', numpy.int32), ('volume_cu_ft', numpy.float32),('flow_cfs', numpy.float32)])

    outArray = numpy.asarray(zip(inArray['OID@'], volume, flow), typeList)

    addMessage('Saving calculations back to feature class')
    arcpy.da.ExtendTable(table, 'OID@', outArray, 'oid', False)

if __name__ == "__main__":
    try:
        outputDir = r'\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality'

        outputGdb = os.path.join(outputDir, 'WQRiskNetwork.gdb')
        linksBase = os.path.join(outputGdb, 'WQConveyance_base')
        nodesBase = os.path.join(outputGdb, 'WQNodes_base')
        inletCatchments = os.path.join(outputGdb, 'InletCatchmentsDetailed')
        flowDir = os.path.join(outputDir, 'flowdir_tau.tif')

        landuseRasterPath = r"\\ws18325\Stormwater D\RiskModel\Surface\LanduseZoningHybrid.tif"
        loadRasterGdb = r"\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality\LoadingEstimates.gdb"

        arcpy.env.overwriteOutput = True
        arcpy.env.snapRaster = flowDir
        arcpy.env.extent = arcpy.Describe(flowDir).Extent
        arcpy.CheckOutExtension("Spatial")

        #Copy the base load network to use later for net load prior to running base load calculation
        linksNet = os.path.join(outputGdb, 'WQConveyance_net')
        nodesNet = os.path.join(outputGdb, 'WQNodes_net')
        #arcpy.Copy_management(linksBase, linksNet)
        #arcpy.Copy_management(nodesBase, nodesNet)

        #CreateLoadRasters(linksNet, nodesNet, inletCatchments, flowDir, landuseRasterPath, loadRasterGdb)
        CalculateBaseFlowAndVolume(linksNet, pd.annualRainfallDepth, 365)
        CalculateBaseFlowAndVolume(nodesNet, pd.annualRainfallDepth, 365)

        Decumulate(linksNet, nodesNet, inletCatchments)

        CalculateScoresAndConcentrations(linksNet, 365)
        CalculateScoresAndConcentrations(nodesNet, 365)
    except Exception as e:
        addMessage("Error: {0}".format(e.message))
    finally:
        addMessage("RAWQ decumulate complete.")